import asyncio
import json
import uuid
from starlette.applications import Starlette
from starlette.routing import Route
from starlette.responses import JSONResponse, StreamingResponse
from starlette.requests import Request
import uvicorn


class TraditionalMCPServer:
    def __init__(self):
        self.active_connections = {}
        self.sessions = {}

    async def handle_sse_connection(self, request: Request):
        session_id = request.headers.get("X-Session-ID")
        if not session_id or session_id not in self.sessions:
            return JSONResponse({"error": "Invalid session"}, status_code=400)

        # 定义SSE流生成器
        async def sse_stream():
            queue = self.active_connections[session_id]
            yield f"data: {json.dumps({
                "jsonrpc":"2.0",
                "method":"notifications/connection",
                "params":{
                    "status":"connected",
                    "session_id":session_id
                }
            })}"
            try:
                while True:
                    # 等待消息队列中的消息
                    message = await asyncio.wait_for(queue.get(), timeout=30)
                    yield f"data:{json.dumps(message)}\n\n"
            except asyncio.TimeoutError:
                yield f"data: {json.dumps({
                    "jsonrpc":"2.0",
                    "method":"notifications/heartbeat",
                    "params":{
                        "timestamp":asyncio.get_event_loop().time()
                    }
                })}"

        return StreamingResponse(
            sse_stream(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Session-ID": session_id,
            },
        )

    async def handle_tool_call(self, request: Request):
        data = await request.json()
        session_id = request.headers.get("X-Session-ID")
        if not session_id or session_id not in self.sessions:
            return JSONResponse({"error": "Invalid session"}, status_code=400)
        tool_name = data["params"]["name"]
        if tool_name == "echo":
            return JSONResponse(
                {
                    "jsonrpc": "2.0",
                    "id": data.get("id"),
                    "result": {
                        "content": [
                            {
                                "type": "text",
                                "text": f"Echo: {data['params']['arguments']['message']}",
                            }
                        ]
                    },
                }
            )
        elif tool_name == "long_running_task":
            duration = data["params"]["arguments"].get("duration", 10)
            return await self.handle_long_running_task(data, session_id, duration)
        else:
            return JSONResponse("error", "Unknown tool", status_code=400)

    async def handle_long_running_task(self, data, session_id, duration):
        async def progress_stream():
            yield f"data: {json.dumps({
                "jsonrpc":"2.0",
                "method":"notification/progress",
                "params":{
                    "type":'begin',
                    "title":"long running task",
                    "message":"Starting task..."
                }
            })}\n\n"

            for i in range(duration):
                await asyncio.sleep(1)
                progress = (i + 1) / duration * 100
                yield f"data: {json.dumps({
                "jsonrpc":"2.0",
                "method":"notification/progress",
                "params":{
                    "type":'update',
                    "title":"long running task",
                    "message":f"progress: {progress:.1f}%",
                    "percentage":progress
                }
            })}\n\n"

            yield f"data: {json.dumps({
                "jsonrpc":"2.0",
                "method":"notification/progress",
                "params":{
                    "type":'end',
                    "title":"long running task",
                    "message":"Task completed"
                }
            })}\n\n"

            yield f"data: {json.dumps({
                "jsonrpc":"2.0",
                "id":data.get('id'),
                "result":{
                    "content":[
                        {
                            "type":"text",
                            "text":f"Long running task completed after {duration} seconds"
                        }
                    ]
                }
            })}\n\n"

        return StreamingResponse(
            progress_stream(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Session-ID": session_id,
            },
        )

    async def handle_tool_list(self, request: Request):
        data = await request.json()
        session_id = request.headers.get("X-Session-ID")
        if not session_id or session_id not in self.sessions:
            return JSONResponse({"error": "Invalid session"}, status_code=400)
        # 工具列表对于无状态和有状态服务器都是一样的
        return JSONResponse(
            {
                "jsonrpc": "2.0",
                "id": data.get("id"),
                "result": {
                    "tools": [
                        {
                            "name": "echo",
                            "description": "Echo back the input (stateless)",
                            "inputSchema": {
                                "type": "object",
                                "properties": {"message": {"type": "string"}},
                                "required": ["message"],
                            },
                        },
                        {
                            "name": "streaming_task",
                            "description": "A task that can stream progress (upgradable to SSE)",
                            "inputSchema": {
                                "type": "object",
                                "properties": {
                                    "duration": {"type": "integer", "default": 5},
                                    "stream": {"type": "boolean", "default": False},
                                },
                            },
                        },
                        {
                            "name": "stateful_counter",
                            "description": "A stateful counter (only works in stateful mode)",
                            "inputSchema": {
                                "type": "object",
                                "properties": {
                                    "action": {
                                        "type": "string",
                                        "enum": ["increment", "decrement", "get"],
                                    }
                                },
                            },
                        },
                    ]
                },
            }
        )

    async def handle_initialize(self, request: Request):
        data = await request.json()
        session_id = str(uuid.uuid4())
        # 保存会话信息
        self.sessions[session_id] = {
            "protocol_version": data.get("protocol_version", "2024-11-05"),
            "capabilities": data.get("capabilities", {}),
            "client_info": data.get("client_info", {}),
            "created_at": asyncio.get_event_loop().time(),
        }
        # 为该会话创建消息队列
        self.active_connections[session_id] = asyncio.Queue()
        return JSONResponse(
            {
                "jsonrpc": "2.0",
                "id": data.get("id"),
                "result": {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {"tools": {}, "resources": {}, "prompts": {}},
                    "serverInfo": {
                        "name": "Traditional MCP Server",
                        "version": "1.0.0",
                        "session_id": session_id,  # 返回会话ID
                    },
                },
            }
        )

    async def handle_shutdown(self, request: Request):
        data = await request.json()
        session_id = request.headers.get("X-Session-ID")
        if session_id in self.sessions:
            del self.sessions[session_id]
        if session_id in self.active_connections:
            del self.active_connections[session_id]
        return JSONResponse({"jsonrpc": "2.0", "id": data.get("id"), "result": {}})


def create_app():
    server = TraditionalMCPServer()

    async def initialize(request: Request):
        return await server.handle_initialize(request)

    async def tool_list(request: Request):
        return await server.handle_tool_list(request)

    async def tool_call(request: Request):
        return await server.handle_tool_call(request)

    async def sse_connection(request: Request):
        return await server.handle_sse_connection(request)

    async def shutdown(request: Request):
        return await server.handle_shutdown(request)

    async def homepage(request: Request):
        return JSONResponse(
            {
                "message": "Traditional HTTP+SSE MCP Server",
                "endpoints": {
                    "/initialize": "POST - initialize MCP Connection",
                    "/tools/list": "POST - list available tools",
                    "/tools/call": "POST - call a tool",
                    "/sse": "POST - establish sse connection",
                    "/shutdown": "POST - shutdown connection",
                },
            }
        )

    return Starlette(
        routes=[
            Route("/", homepage),
            Route("/initialize", initialize, methods=["POST"]),
            Route("/tools/list", tool_list, methods=["POST"]),
            Route("/tools/call", tool_call, methods=["POST"]),
            Route("/sse", sse_connection, methods=["GET"]),
            Route("/shutdown", shutdown, methods=["POST"]),
        ]
    )


if __name__ == "__main__":
    print("starting")
    app = create_app()
    # 1.启动传统的HTTP+SSE的服务器
    # 2.特点
    #    2.1 需要维护长连接
    #    2.2 连接断开后无法恢复
    #    2.3 服务器必须保持高可用性
    #    2.4 只能通过SSE发送服务器消息

    uvicorn.run(app, host="127.0.0.1", port=8000)
